package org.hit.burkun.swalk;

import java.util.HashSet;
import java.util.LinkedList;

import org.hit.burkun.entities.Entity;
import org.hit.burkun.network.EdgeInfo;
import org.hit.burkun.network.SerializableGraph;
import org.hit.burkun.network.Vertex;
import org.hit.burkun.umls.relation.UniformRelation;

public class FeatureExtracter {
	private static int featureNum = 1;

	public static int getFeatureNum() {
		return featureNum;
	}

	private static HashSet<String> doOntologyType = new HashSet<>();
	private static HashSet<String> hpOntologyType = new HashSet<>();
	private static HashSet<String> goOntologyType = new HashSet<>();
	private static HashSet<String> doGeneType = new HashSet<>();
	private static HashSet<String> hpGeneType = new HashSet<>();
	private static HashSet<String> goGeneType = new HashSet<>();
	private static HashSet<String> geneGeneType = new HashSet<>();
	static {
		doOntologyType.add(UniformRelation.TYPE.DO_IS_A.getTag());
		hpOntologyType.add(UniformRelation.TYPE.HP_IS_A.getTag());
		goOntologyType.add(UniformRelation.TYPE.GO_IS_A.getTag());
		goOntologyType.add(UniformRelation.TYPE.GO_PART_OF.getTag());
		goOntologyType.add(UniformRelation.TYPE.GO_REGULTATES.getTag());
		goOntologyType.add(UniformRelation.TYPE.GO_HAPPENS_DURING.getTag());
		goOntologyType.add(UniformRelation.TYPE.GO_OCCURS_IN.getTag());

		doGeneType.add(UniformRelation.TYPE.DO_SIDD_MFI_MESH_GENE.getTag());
		doGeneType.add(UniformRelation.TYPE.DO_SIDD_MFI_OMIM_GENE.getTag());
		doGeneType.add(UniformRelation.TYPE.DO_SIDD_MFS_MESH_GENE.getTag());
		doGeneType.add(UniformRelation.TYPE.DO_SIDD_MFR_MESH_GENE.getTag());
		doGeneType.add(UniformRelation.TYPE.DO_SIDD_MFR_OMIM_GENE.getTag());
		doGeneType.add(UniformRelation.TYPE.DO_XREF_OMIM_GENE.getTag());
		doGeneType.add(UniformRelation.TYPE.DISEASE_CTD_GENE_CURATED.getTag());
		doGeneType.add(UniformRelation.TYPE.DISEASE_CLINVAR_GENE.getTag());

		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_EXP.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_IPI.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_IDA.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_IMP.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_IGI.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_IEP.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_IKR.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_ISS.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_IBA.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_IRD.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_TAS.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_IC.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_ND.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_IEA.getTag());
		goGeneType.add(UniformRelation.TYPE.GO_ANNOTATION_GENE_NAS.getTag());

		hpGeneType.add(UniformRelation.TYPE.HP_ANNOTATION_GENE_ICE.getTag());
		hpGeneType.add(UniformRelation.TYPE.HP_ANNOTATION_GENE_PCS.getTag());
		hpGeneType.add(UniformRelation.TYPE.HP_ANNOTATION_GENE_TAS.getTag());
		hpGeneType.add(UniformRelation.TYPE.HP_ANNOTATION_GENE_IEA.getTag());
		hpGeneType.add(UniformRelation.TYPE.HP_FILE_GENE.getTag());
		hpGeneType.add(UniformRelation.TYPE.DISEASE_CTD_GENE_CURATED.getTag());
		hpGeneType.add(UniformRelation.TYPE.DISEASE_CLINVAR_GENE.getTag());

		geneGeneType.add(UniformRelation.TYPE.STRING_GENE.getTag());
	}

	public static double[] extractFeaturesVersion8Score(Vertex from, Vertex to,
			SerializableGraph sg) {
		featureNum = 8;
		LinkedList<EdgeInfo> edges = sg.getNeighborsByVertex(from);
		/*
		 * 0 : do-ontology 1 : hp-ontology 2 : go-ontology 3 : do-gene 4 :
		 * go-gene 5 : hp-gene 6 : gene-gene 8 : comment nei 9 : degree
		 */
		double[] features = new double[featureNum];
		for (EdgeInfo e : edges) {
			if (e.getVertex().getUid() == to.getUid()) {
				if (e.isInRelType(doOntologyType)) {
					features[0] = e.getScore();;
				}
				if (e.isInRelType(hpOntologyType)) {
					features[1] = e.getScore();;
				}
				if (e.isInRelType(goOntologyType)) {
					features[2] = e.getScore();;
				}
				if (e.isInRelType(doGeneType)
						&& (from.isInNodeTypes(Entity.TYPE.DISEASE_ONTOLOGY
								.getTag()) || e.getVertex().isInNodeTypes(
								Entity.TYPE.DISEASE_ONTOLOGY.getTag()))) {
					features[3] = e.getScore();;
				}
				if (e.isInRelType(goGeneType)) {
					features[4] = e.getScore();;
				}
				if (e.isInRelType(hpGeneType)
						&& (from.isInNodeTypes(Entity.TYPE.HP_ONTOLOGY.getTag()) || e
								.getVertex().isInNodeTypes(
										Entity.TYPE.HP_ONTOLOGY.getTag()))) {
					features[5] = e.getScore();;
				}
				if (e.isInRelType(geneGeneType)) {
					features[6] = e.getScore();;
				}
				break;
			}
		}
		features[7] = 1;
		return features;
	}
	
	public static double[] extractFeaturesVersion10f(Vertex from, Vertex to,
			SerializableGraph sg) {
		featureNum = 10;
		LinkedList<EdgeInfo> edges = sg.getNeighborsByVertex(from);
		/*
		 * 0 : do-ontology 1 : hp-ontology 2 : go-ontology 3 : do-gene 4 :
		 * go-gene 5 : hp-gene 6 : gene-gene 8 : comment nei 9 : degree
		 */
		double[] features = new double[featureNum];
		for (EdgeInfo e : edges) {
			if (e.getVertex().getUid() == to.getUid()) {
				if (e.isInRelType(doOntologyType)) {
					features[0] = 1;
				}
				if (e.isInRelType(hpOntologyType)) {
					features[1] = 1;
				}
				if (e.isInRelType(goOntologyType)) {
					features[2] = 1;
				}
				if (e.isInRelType(doGeneType)
						&& (from.isInNodeTypes(Entity.TYPE.DISEASE_ONTOLOGY
								.getTag()) || e.getVertex().isInNodeTypes(
								Entity.TYPE.DISEASE_ONTOLOGY.getTag()))) {
					features[3] = 1;
				}
				if (e.isInRelType(goGeneType)) {
					features[4] = 1;
				}
				if (e.isInRelType(hpGeneType)
						&& (from.isInNodeTypes(Entity.TYPE.HP_ONTOLOGY.getTag()) || e
								.getVertex().isInNodeTypes(
										Entity.TYPE.HP_ONTOLOGY.getTag()))) {
					features[5] = 1;
				}
				if (e.isInRelType(geneGeneType)) {
					features[6] = 1;
				}
				features[8] = e.getScore();
				break;
			}
		}
		features[7] = 1;
		LinkedList<EdgeInfo> fromEdge = sg.getNeighborsByVertex(from);
		LinkedList<EdgeInfo> toEdge = sg.getNeighborsByVertex(to);
		HashSet<EdgeInfo> fromSet = new HashSet<>(fromEdge);
		HashSet<EdgeInfo> toSet = new HashSet<>(toEdge);
		fromSet.retainAll(toSet);
		features[9] = fromSet.size();
		if (features[9] != 0) {
			toSet.addAll(fromEdge);
			features[9] /= toSet.size();
		}
		return features;
	}


	public static double[] extractFeaturesVersion9PlusCn(Vertex from, Vertex to,
			SerializableGraph sg) {
		featureNum = 9;
		LinkedList<EdgeInfo> edges = sg.getNeighborsByVertex(from);
		/*
		 * 0 : do-ontology 1 : hp-ontology 2 : go-ontology 3 : do-gene 4 :
		 * go-gene 5 : hp-gene 6 : gene-gene 8 : comment nei 9 : degree
		 */
		double[] features = new double[featureNum];
		for (EdgeInfo e : edges) {
			if (e.getVertex().getUid() == to.getUid()) {
				if (e.isInRelType(doOntologyType)) {
					features[0] = e.getScore();;
				}
				if (e.isInRelType(hpOntologyType)) {
					features[1] = e.getScore();;
				}
				if (e.isInRelType(goOntologyType)) {
					features[2] = e.getScore();;
				}
				if (e.isInRelType(doGeneType)
						&& (from.isInNodeTypes(Entity.TYPE.DISEASE_ONTOLOGY
								.getTag()) || e.getVertex().isInNodeTypes(
								Entity.TYPE.DISEASE_ONTOLOGY.getTag()))) {
					features[3] = e.getScore();;
				}
				if (e.isInRelType(goGeneType)) {
					features[4] = e.getScore();;
				}
				if (e.isInRelType(hpGeneType)
						&& (from.isInNodeTypes(Entity.TYPE.HP_ONTOLOGY.getTag()) || e
								.getVertex().isInNodeTypes(
										Entity.TYPE.HP_ONTOLOGY.getTag()))) {
					features[5] = e.getScore();;
				}
				if (e.isInRelType(geneGeneType)) {
					features[6] = e.getScore();;
				}
				break;
			}
		}
		features[7] = 1;
		LinkedList<EdgeInfo> fromEdge = sg.getNeighborsByVertex(from);
		LinkedList<EdgeInfo> toEdge = sg.getNeighborsByVertex(to);
		HashSet<EdgeInfo> fromSet = new HashSet<>(fromEdge);
		HashSet<EdgeInfo> toSet = new HashSet<>(toEdge);
		fromSet.retainAll(toSet);
		features[8] = fromSet.size();
		return features;
	}

	// 7个 socore
	public static double[] extractFeaturesVersion7Score(Vertex from, Vertex to,
			SerializableGraph sg) {
		featureNum = 7;
		LinkedList<EdgeInfo> edges = sg.getNeighborsByVertex(from);
		/*
		 * 0 : do-ontology 1 : hp-ontology 2 : go-ontology 3 : do-gene 4 :
		 * go-gene 5 : hp-gene 6 : gene-gene
		 */
		double[] features = new double[featureNum];
		for (EdgeInfo e : edges) {
			if (e.getVertex().getUid() == to.getUid()) {
				if (e.isInRelType(doOntologyType)) {
					features[0] = e.getScore();
				}
				if (e.isInRelType(hpOntologyType)) {
					features[1] = e.getScore();
				}
				if (e.isInRelType(goOntologyType)) {
					features[2] = e.getScore();
				}
				if (e.isInRelType(doGeneType)
						&& (from.isInNodeTypes(Entity.TYPE.DISEASE_ONTOLOGY
								.getTag()) || e.getVertex().isInNodeTypes(
								Entity.TYPE.DISEASE_ONTOLOGY.getTag()))) {
					features[3] = e.getScore();
				}
				if (e.isInRelType(goGeneType)) {
					features[4] = e.getScore();
				}
				if (e.isInRelType(hpGeneType)
						&& (from.isInNodeTypes(Entity.TYPE.HP_ONTOLOGY.getTag()) || e
								.getVertex().isInNodeTypes(
										Entity.TYPE.HP_ONTOLOGY.getTag()))) {
					features[5] = e.getScore();
				}
				if (e.isInRelType(geneGeneType)) {
					features[6] = e.getScore();
				}
				break;
			}
		}

		return features;
	}

	// 加入常数项
	public static double[] extractFeaturesVersion6(Vertex from, Vertex to,
			SerializableGraph sg) {
		featureNum = 8;
		LinkedList<EdgeInfo> edges = sg.getNeighborsByVertex(from);
		/*
		 * 0 : do-ontology 1 : hp-ontology 2 : go-ontology 3 : do-gene 4 :
		 * go-gene 5 : hp-gene 6 : gene-gene
		 */
		double[] features = new double[featureNum];
		for (EdgeInfo e : edges) {
			if (e.getVertex().getUid() == to.getUid()) {
				if (e.isInRelType(doOntologyType)) {
					features[0] = e.getScore();
				}
				if (e.isInRelType(hpOntologyType)) {
					features[1] = e.getScore();
				}
				if (e.isInRelType(goOntologyType)) {
					features[2] = e.getScore();
				}
				if (e.isInRelType(doGeneType)
						&& (from.isInNodeTypes(Entity.TYPE.DISEASE_ONTOLOGY
								.getTag()) || e.getVertex().isInNodeTypes(
								Entity.TYPE.DISEASE_ONTOLOGY.getTag()))) {
					features[3] = e.getScore();
				}
				if (e.isInRelType(goGeneType)) {
					features[4] = e.getScore();
				}
				if (e.isInRelType(hpGeneType)
						&& (from.isInNodeTypes(Entity.TYPE.HP_ONTOLOGY.getTag()) || e
								.getVertex().isInNodeTypes(
										Entity.TYPE.HP_ONTOLOGY.getTag()))) {
					features[5] = e.getScore();
				}
				if (e.isInRelType(geneGeneType)) {
					features[6] = e.getScore();
				}
				break;
			}
		}
		features[7] = 1;
		return features;
	}
	//三阶邻居
	public static double[] extractFeaturesVersion10WithScoref(Vertex from, Vertex to,
			SerializableGraph sg) {
		featureNum = 10;
		LinkedList<EdgeInfo> edges = sg.getNeighborsByVertex(from);
		/*
		 * 0 : do-ontology 1 : hp-ontology 2 : go-ontology 3 : do-gene 4 :
		 * go-gene 5 : hp-gene 6 : gene-gene 8 : comment nei 9 : degree
		 */
		double[] features = new double[featureNum];
		for (EdgeInfo e : edges) {
			if (e.getVertex().getUid() == to.getUid()) {
				if (e.isInRelType(doOntologyType)) {
					features[0] = e.getScore();
				}
				if (e.isInRelType(hpOntologyType)) {
					features[1] = e.getScore();
				}
				if (e.isInRelType(goOntologyType)) {
					features[2] = e.getScore();
				}
				if (e.isInRelType(doGeneType)
						&& (from.isInNodeTypes(Entity.TYPE.DISEASE_ONTOLOGY
								.getTag()) || e.getVertex().isInNodeTypes(
								Entity.TYPE.DISEASE_ONTOLOGY.getTag()))) {
					features[3] = e.getScore();
				}
				if (e.isInRelType(goGeneType)) {
					features[4] = e.getScore();
				}
				if (e.isInRelType(hpGeneType)
						&& (from.isInNodeTypes(Entity.TYPE.HP_ONTOLOGY.getTag()) || e
								.getVertex().isInNodeTypes(
										Entity.TYPE.HP_ONTOLOGY.getTag()))) {
					features[5] = e.getScore();
				}
				if (e.isInRelType(geneGeneType)) {
					features[6] = e.getScore();
				}
				break;
			}
		}
		features[7] = 1;
		LinkedList<EdgeInfo> fromEdge = sg.getNeighborsByVertex(from);
		LinkedList<EdgeInfo> toEdge = sg.getNeighborsByVertex(to);
		HashSet<EdgeInfo> fromSet = new HashSet<>(fromEdge);
		HashSet<EdgeInfo> toSet = new HashSet<>(toEdge);
		HashSet<EdgeInfo> fromSet2 = new HashSet<>();
		HashSet<EdgeInfo> toSet2 = new HashSet<>();
		//2neibor
		for(EdgeInfo e : fromSet){
			LinkedList<EdgeInfo> edges2 = sg.getNeighborsByVertex(e.getVertex());
			fromSet2.addAll(edges2);
		}
		for(EdgeInfo e : toSet){
			LinkedList<EdgeInfo> edges2 = sg.getNeighborsByVertex(e.getVertex());
			toSet2.addAll(edges2);
		}
		fromSet.retainAll(toSet);
		features[8] = fromSet.size();
		fromSet2.retainAll(toSet2);
		features[9] = fromSet2.size();
		return features;
	}
	public static double[] extractFeaturesAllOne() {
		featureNum = 1;
		double[] fs = new double[1];
		fs[0] = 1;
		return fs;
	}

	public static void setFeatureNum(int num) {
		featureNum = num;
	}
}
